import pandas as pd
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import torch
import math
import sys
from torch.utils.data import DataLoader,TensorDataset
from demo import model
from model import NeuralNet
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device=', device)
df_ = pd.read_csv('diabetes.csv')
df = df_.copy()
data = df.iloc[:, :-1].to_numpy()
label = df.iloc[:, -1:].to_numpy()
x = torch.tensor(data, dtype=torch.float32).to(device)
target = torch.tensor(label, dtype=torch.long).to(device)
length = x.shape[0]
train_len = int(length * 0.8)
train_data = x[:train_len, :]
train_label = target[:train_len, :]
val_data = x[train_len:, :]
val_label = target[train_len:, :]

train_dataset = TensorDataset(train_data, train_label)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)

val_dataset = TensorDataset(val_data, val_label)
val_length = val_data.shape[0]
val_loader = DataLoader(val_dataset, batch_size=val_length, shuffle=True, num_workers=0)
val_item, label_item = next(iter(val_loader))

net = model(8, 2)
# net = NeuralNet(8, 2)
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002, weight_decay=1e-3)

epochs = 10
save_path = './Net.pth'
best_acc = 0.0
train_steps = len(train_loader)
for epoch in range(epochs):
    # train
    net.train()
    running_loss = 0.0
    for step, data in enumerate(train_loader):
        images, labels = data
        optimizer.zero_grad()
        outputs = net(images.to(device))
        loss = loss_function(outputs, labels.squeeze(dim=1).to(device))
        loss.backward()
        optimizer.step()
        # print statistics
        running_loss += loss.item()
        # if step+1 == train_steps:  # print every 500 mini-batches
        with torch.no_grad():
            outputs = net(val_item)  # [batch, 10]
            predict_y = torch.max(outputs, dim=1)[1]
            accuracy = torch.eq(predict_y, label_item).sum().item() / label_item.size(0)
            print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                  (epoch + 1, step + 1, running_loss, accuracy))
            running_loss = 0.0
print('Finished Training')
save_path = './net.pth'
torch.save(net.state_dict(), save_path)
